from keras.utils import to_categorical
from keras import models, layers, regularizers
from keras.models import load_model
from keras.optimizers import RMSprop
from keras.datasets import mnist


def load(model_name):
    return load_model(model_name)
    
def mnist_init():
    """ 加载 mnist 数据集 """
    (t_images, t_labels), (te_images, te_labels) = mnist.load_data()

    global train_images 
    train_images = t_images.reshape((60000, 28*28)).astype('float')
    global test_images 
    test_images = te_images.reshape((10000, 28*28)).astype('float')
    global train_labels 
    train_labels = to_categorical(t_labels)
    global test_labels 
    test_labels = to_categorical(te_labels)
    
def build_network(shape):
    """
    构造神经网络
    @parem shape
    @return network

    shape 示例：[64,128,128]
        代表有三层中间层，每层神经元的个数分别是64,128,128
    """
    
    network = models.Sequential()

    # 输入层
    network.add(layers.Dense(
        units= 128, activation= "relu", input_shape=(28*28, ),
        kernel_regularizer= regularizers.l1(0.0001)))
    
    network.add(layers.Dropout(0.01))
        
    # 中间层
    for i in shape:
        network.add(layers.Dense(
            units= i, activation= "relu",
            kernel_regularizer=regularizers.l1(0.0001)))
        
        network.add(layers.Dropout(0.01))
        
    # 输出层
    network.add(layers.Dense(units=10, activation='softmax'))

    return network

def train(network,lrate,ep):
    """
    训练神经网络
    @parem learning_rate, epochs
    @return network

    learning_rate 学习率，即梯度下降每步的步长因数
    epochs 训练的轮数
    
    """
    # 编译步骤
    network.compile(
            optimizer=RMSprop(lr= lrate),
            loss='categorical_crossentropy', metrics=['accuracy'])

    # 训练网络，用fit函数, epochs表示训练多少个回合， batch_size表示每次训练给多大的数据
    network.fit(train_images, train_labels, epochs= ep, batch_size=128, verbose=2)
    return network

def test_model(network):
    """
    评价模型
    @parem network
    @return [loss, test_accuracy]
    
    loss 是损失函数值
    test_accuracy 是成功率
    """

    return network.evaluate(test_images, test_labels)

def __init__(model_name):
    # 初始化
    return load(model_name)

if __name__ == "__main__":

    mnist_init()
    # 构造训练用神经网络
    network = build_network([256,128])
    try:
        network = train(network, 0.001, 20)
    except EOFError:
        network.save("EOFError.h5")
    else:
        network.save("accident.h5")
    
    network.save("higher.h5")
    
    #test_network = __init__()

    test_loss, test_accuracy = test_model(network)
    print("test_loss:", test_loss, "    test_accuracy:", test_accuracy)


















